from pprint import pprint
from paddlenlp import Taskflow
# # 定义实体关系抽取的schema
# schema = ['姓名', '出生日期', '电话']
# ie = Taskflow('information_extraction', schema=schema)

from docx import Document
from docx.shared import Inches

import datetime
import os
import fitz  # fitz就是pip install PyMuPDF
import cv2
import shutil
import numpy as np
import pandas as pd
from tqdm import tqdm


def get_paragraphs_text(path):
    """
    获取所有段落的文本
    :param path: word路径
    :return: list类型，如：
        ['Test', 'hello world', ...]
    """
    document = Document(path) 
    # 有的简历是表格式样的，因此，不仅需要提取正文，还要提取表格
    col_keys = [] # 获取列名
    col_values = [] # 获取列值
    index_num = 0
    # 表格提取中，需要添加一个去重机制
    fore_str = ""
    cell_text = ""
    for table in document.tables:
        for row_index,row in enumerate(table.rows):
            for col_index,cell in enumerate(row.cells):
                if fore_str != cell.text:
                    if index_num % 2==0:
                        col_keys.append(cell.text)
                    else:
                        col_values.append(cell.text)
                    fore_str = cell.text
                    index_num +=1
                    cell_text += cell.text + '\n'
    # 提取正文文本
    paragraphs_text = ""
    for paragraph in document.paragraphs:
        # 拼接一个list,包括段落的结构和内容
        paragraphs_text += paragraph.text + "\n"
    return cell_text, paragraphs_text


def pyMuPDF_fitz(pdfPath, imagePath):
    startTime_pdf2img = datetime.datetime.now()  # 开始时间

    print("imagePath=" + imagePath)
    pdfDoc = fitz.open(pdfPath)
    for pg in range(pdfDoc.page_count):
        page = pdfDoc[pg]
        rotate = int(0)
        # 每个尺寸的缩放系数为4，这将为我们生成分辨率提高4的图像。
        # 此处若是不做设置，默认图片大小为：792X612, dpi=96
        zoom_x = 4  # (1.33333333-->1056x816)   (2-->1584x1224)
        zoom_y = 4
        mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
        pix = page.get_pixmap(matrix=mat, alpha=False)

        if not os.path.exists(imagePath):  # 判断存放图片的文件夹是否存在
            os.makedirs(imagePath)  # 若图片文件夹不存在就创建

        pix.save(imagePath + '/' + 'images_%s.jpeg' % pg)  # 将图片写入指定的文件夹内

    endTime_pdf2img = datetime.datetime.now()  # 结束时间
    print('pdf转图片时间=', (endTime_pdf2img - startTime_pdf2img).seconds)


# 选定一些HR关注的基本信息
schema = ['姓名', '出生日期', '电话', '性别', '最高学历', '籍贯', '政治面貌', '毕业院校', '学位', '毕业时间', '工作时间',"项目经验","实习","能力","技能","意向岗位","岗位","项目介绍","项目内容","项目名称"]

# schema = ['姓名', '出生日期', '电话', '性别', '最高学历', '籍贯', '政治面貌', '毕业院校', '学位', '毕业时间', '工作时间']
ie = Taskflow('information_extraction', schema=schema)

def get_info(cell_text, paragraphs_text):
    # 将整理后的抽取结果返回为字典
    schema_dict = {}
    # 抽取简历信息
    a = ie(cell_text, paragraphs_text)
    for i in schema:
        if i in a[0]:
            schema_dict[i] = a[0][i][0]['text']
            # 查看抽取信息
            # print(a[0][i][0]['text'])
        else:
            schema_dict[i] = ''
    return schema_dict



def get_word(path,max_idx=100):
    filenames = os.listdir(path)
    result = []
    idx=0
    for filename in tqdm(filenames):
        idx+=1
        # if idx>=max_idx:
        #     break
        if idx<=2:
            continue
        cell_text, paragraphs_text = get_paragraphs_text(os.path.join(path,filename))
        res = get_info(cell_text, paragraphs_text)
        res['filename'] = filename
        result.append(res)
        result_pd = pd.DataFrame(result)
        # result_pd.to_excel(f'简历信息_data_{idx}.xlsx')
#         result_pd.to_excel(f'out_excel/简历信息_data_{idx}.xlsx')
        # result_pd.to_excel(f'/kaggle/working/out_excel/简历信息_data_{idx}.xlsx')
        result_pd.to_excel(f'{out_dir}/简历信息_data_{idx}.xlsx')
    return result



def get_pic_info(path):
    # 将整理后的抽取结果返回为字典
    schema_dict = {}
    if os.path.splitext(path)[-1]=='.pdf':
        pdfDoc = fitz.open(path)
        for pg in range(pdfDoc.page_count):
            page = pdfDoc[pg]
            rotate = int(0)
            zoom_x = 4  # (1.33333333-->1056x816)   (2-->1584x1224)
            zoom_y = 4
            mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
            pix = page.get_pixmap(matrix=mat, alpha=False)
            # 保存过渡图片
            temp_path = pix.save('temp.jpeg')
            # 抽取过渡图片中的简历信息
            a = ie({"doc":'temp.jpeg'})
            if pg==0:
                for i in schema:
                    if i in a[0]:
                        schema_dict[i] = a[0][i][0]['text']
                    else:
                        schema_dict[i] = ''
            else:
                for i in schema:
                    if i in a[0]:
                        schema_dict[i] = a[0][i][0]['text']
    elif os.path.splitext(path)[-1]=='.jpeg':
        # 抽取图片中的简历信息
        a = ie({"doc":path})
        for i in schema:
            if i in a[0]:
                schema_dict[i] = a[0][i][0]['text']
            else:
                schema_dict[i] = ''
    else:
        print('图片信息抽取只支持pdf和jpeg格式，请调整后重试。')
    return schema_dict


def get_pics(path):
    filenames = os.listdir(path)
    result = []
    for filename in tqdm(filenames):
        res = get_pic_info(os.path.join(path,filename))
        res['文件名'] = filename
        result.append(res)
    return result


out_dir="/kaggle/working/out_excel"
os.makedirs(out_dir,exist_ok=True)


docx_dir="/kaggle/input/resume-train-20200121/train_20200121/resume_train_20200121/docx"
result = get_word(docx_dir)

result_pd = pd.DataFrame(result)
result_pd.to_excel('简历信息_data.xlsx')

